#!/usr/bin/env python
import numpy as np
import rospy
import tf2_ros
from tf.transformations import quaternion_matrix, quaternion_from_matrix
from std_msgs.msg import String
from geometry_msgs.msg import PoseStamped,TransformStamped
from geometry_msgs.msg import Pose
from scipy.spatial.transform import Rotation as R



# 定义旋转矩阵函数
def get_transform(target_frame, source_frame):
    """
    获取指定两个坐标系之间的变换
    """


    try:
        transform = tf_buffer.lookup_transform(target_frame, source_frame, rospy.Time(0))
        translation = np.array([transform.transform.translation.x,
                                transform.transform.translation.y,
                                transform.transform.translation.z])
        rotation = np.array([transform.transform.rotation.x,
                            transform.transform.rotation.y,
                            transform.transform.rotation.z,
                            transform.transform.rotation.w])
        R = quaternion_matrix(rotation)[:3, :3]
        return translation, R
    except tf2_ros.LookupException as e:
        rospy.logwarn("Transform not available: %s", str(e))
    except tf2_ros.ConnectivityException as e:
        rospy.logerr("Connectivity issue: %s", str(e))
    except tf2_ros.ExtrapolationException as e:
        rospy.logerr("Extrapolation issue: %s", str(e))
    return None, None



def compute_grasp_position(P_object, R_object, shape, dimensions=None):
    """
    计算夹爪接触点和姿态
    Parameters:
    - P_object: 物体位置 (3x1 array)
    - R_object: 物体姿态 (3x3 rotation matrix)
    - shape: 物体形状 ('cylinder' or 'cuboid')
    - dimensions: 物体尺寸 (tuple), cylinder: (r, h), cuboid: (l, w, h)
    - grasp_distance: 抓取点到物体的距离
    Returns:
    - P_grasp: 抓取点的位置
    - grasp_orientation: 抓取姿态 (3x3 rotation matrix)
    """
    if shape == 'cylinder':
        r, h = dimensions
        # 设置抓取点在物体X轴正方向上


        obj_x_dir = R_object[:, 0]
        obj_y_dir = R_object[:, 1]
        obj_z_dir = R_object[:, 2]

        tool_x_dir = -obj_z_dir
        tool_x_dir = tool_x_dir / np.linalg.norm(tool_x_dir)
        tool_y_candidate = np.array([0.0,-1.0,0.0])

        dot_xy = np.abs(np.dot(tool_x_dir,tool_y_candidate))

        if dot_xy > 0.99:
            tool_y_candidate = np.array([0.0,0.0,-1.0])
        
        v_perp = tool_y_candidate - np.dot(tool_y_candidate,tool_x_dir) * tool_x_dir
        norm_vp = np.linalg.norm(v_perp)
        
        if norm_vp < 1e-8:
            tool_y_candidate = np.array([0.0,0.0,1.0])
            v_perp = tool_y_candidate - np.dot(tool_y_candidate,tool_x_dir) * tool_x_dir
            norm_vp = np.linalg.norm(v_perp)

        tool_y_dir = v_perp / norm_vp
        tool_z_dir = np.cross(tool_x_dir,tool_y_dir)
        tool_z_dir = tool_z_dir / np.linalg.norm(tool_z_dir)

        grasp_orientation = np.column_stack((tool_x_dir, tool_y_dir, tool_z_dir))
        grasp_poses =[]
        grasp_poses.append({
                "position": P_object + grasp_orientation @ np.array([0, r, 0.063])  ,# 沿 y 轴正方向移动  z轴前靠6cm
                "orientation": grasp_orientation
            })

 

        ''''''
        br = tf2_ros.TransformBroadcaster()
        t = TransformStamped()
        t.header.stamp = rospy.Time.now()
        t.header.frame_id = "base_link"
        t.child_frame_id = "target_tool_link"

        t.transform.translation.x=grasp_poses[0]["position"][0]
        t.transform.translation.y=grasp_poses[0]["position"][1]
        t.transform.translation.z=grasp_poses[0]["position"][2]

        R_object_orientation = np.eye(4)
        R_object_orientation[:3, :3] = grasp_poses[0]["orientation"]
        t.transform.rotation.x = quaternion_from_matrix(R_object_orientation)[0]
        t.transform.rotation.y = quaternion_from_matrix(R_object_orientation)[1]
        t.transform.rotation.z = quaternion_from_matrix(R_object_orientation)[2]
        t.transform.rotation.w = quaternion_from_matrix(R_object_orientation)[3]
        br.sendTransform(t)



    elif shape == 'cuboid':
        l, w, h = dimensions
        # 设置抓取点在物体X轴正方向上
        P_contact = np.array([l / 2 + grasp_distance, 0, 0])
        # 计算抓取点在全局坐标系下的位置
        P_grasp = P_object + np.dot(R_object, P_contact)
        # 构建抓取姿态
        x_axis = np.array([1, 0, 0])  # X轴正方向
        z_axis = R_object[:, 2]  # 物体的Z轴方向
        y_axis = np.cross(z_axis, x_axis)
        y_axis /= np.linalg.norm(y_axis)  # 单位化
        z_axis = np.cross(x_axis, y_axis)  # 调整Z轴
        grasp_orientation = np.column_stack((x_axis, y_axis, z_axis))
    
    else:
        raise ValueError("Unsupported shape: choose 'cylinder' or 'cuboid'")



    return grasp_poses[0]["position"],grasp_poses[0]["orientation"]

def apply_fixed_transform(P_object, R_object, transform_params):
    """
    应用固定的变换参数
    """
    qw = transform_params["qw"]
    qx = transform_params["qx"]
    qy = transform_params["qy"]
    qz = transform_params["qz"]
    x = transform_params["x"]
    y = transform_params["y"]
    z = transform_params["z"]

    fixed_rotation = np.array([qx, qy, qz, qw])
    R_fixed = quaternion_matrix(fixed_rotation)[:3, :3]
    P_fixed = np.array([x, y, z])

    P_object_transformed = np.dot(R_fixed, P_object) + P_fixed
    R_object_transformed = np.dot(R_fixed, R_object)

    return P_object_transformed, R_object_transformed


def process_pose_data(msg):
    """
    处理从PoseStamped消息中接收到的pose数据并计算抓取位姿
    """
    # 从PoseStamped消息中提取物体的位置和四元数姿态
    P_object = np.array([msg.pose.position.x, msg.pose.position.y, msg.pose.position.z])
    quaternion = [msg.pose.orientation.x, msg.pose.orientation.y, msg.pose.orientation.z, msg.pose.orientation.w]

    # 将四元数转换为旋转矩阵
    R_object = quaternion_matrix(quaternion)[:3, :3]



    # 定义物体形状和尺寸（假设是圆柱体）
    shape = 'cylinder'
    dimensions = [0.070, 0.106]  # 直径和高度

    # 计算抓取点位置和抓取姿态
    P_grasp, grasp_orientation = compute_grasp_position(P_object, R_object, shape, dimensions)
    ''''''
    br = tf2_ros.TransformBroadcaster()
    t = TransformStamped()
    t.header.stamp = rospy.Time.now()
    t.header.frame_id = "base_link"
    t.child_frame_id = "obj_link"

    t.transform.translation.x=P_object[0]
    t.transform.translation.y=P_object[1]
    t.transform.translation.z=P_object[2]

    R_object_orientation = np.eye(4)
    R_object_orientation[:3, :3] = R_object
    t.transform.rotation.x = quaternion_from_matrix(R_object_orientation)[0]
    t.transform.rotation.y = quaternion_from_matrix(R_object_orientation)[1]
    t.transform.rotation.z = quaternion_from_matrix(R_object_orientation)[2]
    t.transform.rotation.w = quaternion_from_matrix(R_object_orientation)[3]
    br.sendTransform(t)
    
    

    # 将抓取姿态扩展为 4x4 齐次变换矩阵
    grasp_orientation_homogeneous = np.eye(4)
    grasp_orientation_homogeneous[:3, :3] = grasp_orientation

    # 将抓取姿态转换为四元数
    grasp_quaternion = quaternion_from_matrix(grasp_orientation_homogeneous)

    # 构建抓取位姿消息
    grasp_msg = Pose()
    grasp_msg.position.x = P_grasp[0]
    grasp_msg.position.y = P_grasp[1]
    grasp_msg.position.z = P_grasp[2]
    grasp_msg.orientation.x = grasp_quaternion[0]
    grasp_msg.orientation.y = grasp_quaternion[1]
    grasp_msg.orientation.z = grasp_quaternion[2]
    grasp_msg.orientation.w = grasp_quaternion[3]
    '''
    {
        "position": {
            "x": P_grasp[0],
            "y": P_grasp[1],
            "z": P_grasp[2],
        },
        "orientation": {
            "x": grasp_quaternion[0],
            "y": grasp_quaternion[1],
            "z": grasp_quaternion[2],
            "w": grasp_quaternion[3],
        }
    }
    '''
    # 发布抓取位姿
    grasp_publisher.publish(grasp_msg)
    rospy.loginfo(f"Grasp Pose Published: {grasp_msg}")




if __name__ == "__main__":
    rospy.init_node("grasp_calculation_node")
    grasp_publisher = rospy.Publisher("/right_grasp_pose", Pose, queue_size=1)
    tf_buffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tf_buffer)
    # rospy.sleep(1.0)  # 等待缓冲区初始化

    rospy.Subscriber("/obj_pose", PoseStamped, process_pose_data, queue_size=1)

    rospy.spin()
